test_algo_simple_graph = """{
    "properties": {},
    "inports": {},
    "outports": {},
    "processes": {
        "start": {
            "component": "start",
            "metadata": {
                "x": -140,
                "y": -60,
                "width": 60,
                "height": 60,
                "label": "开始"
            }
        },
        "end": {
            "component": "end",
            "metadata": {
                "x": 400,
                "y": -60,
                "width": 60,
                "height": 60,
                "label": "结束"
            }
        },
        "b342a8bd-21db-47be-89a1-0db7ab1798fb": {
            "component": "add",
            "metadata": {
                "x": -10,
                "y": -50,
                "width": 120,
                "height": 40,
                "label": "累加"
            }
        },
        "9fbad8cb-110e-4cd3-97ee-cbe7919b8753": {
            "component": "multiple",
            "metadata": {
                "x": 190,
                "y": -50,
                "width": 120,
                "height": 40,
                "label": "乘法"
            }
        }
    },
    "connections": [
        {
            "src": {
                "process": "start",
                "port": "out"
            },
            "tgt": {
                "process": "b342a8bd-21db-47be-89a1-0db7ab1798fb",
                "port": "in"
            }
        },
        {
            "src": {
                "process": "b342a8bd-21db-47be-89a1-0db7ab1798fb",
                "port": "out"
            },
            "tgt": {
                "process": "9fbad8cb-110e-4cd3-97ee-cbe7919b8753",
                "port": "in"
            }
        },
        {
            "src": {
                "process": "9fbad8cb-110e-4cd3-97ee-cbe7919b8753",
                "port": "out"
            },
            "tgt": {
                "process": "end",
                "port": "in"
            }
        }
    ]
}"""

import logging
logging.basicConfig(level=logging.INFO)
from calc_manager.operator.operator import * 
from calc_manager.operator.repo import * 
from typing import Optional

class FakeOperatorRepository(OperatorRepository):
    def __init__(self):
        sample_add = {
            "operator_name": "add",
            "operator_path": "operators_sample/sample_add/add.py",
            "operator_func_entrypoint": "add",
            "operator_desc": "This is a test operator",
        }
        sample_mul = {
            "operator_name": "multiple",
            "operator_path": "operators_sample/sample_multiple/mul.py",
            "operator_func_entrypoint": "mul",
            "operator_desc": "This is a test operator",
        }
        self.repo = {"add": sample_add, "multiple": sample_mul}
        

    def get_op_by_id(self, op_id: str) -> Optional[OperatorInfo]:
        info = self.repo[op_id]
        return OperatorInfo(**info)

        
from calc_manager.algorithm.algo import Algorithm, AlgorithmInfo, OperatorComponentToOperatorInfoAdapter
from calc_manager.algorithm.graph import AlgoGraph
from calc_manager.algorithm.param import ParameterMappingService
from calc_manager.algorithm.types import ParameterMapping
from calc_manager.executor.exec import AlgoTaskFactory, build_algo_tasks, DefaultExecutor

algo_graph = AlgoGraph(test_algo_simple_graph, "start", "end")
algo_info = AlgorithmInfo(algo_name="test", graph_json=test_algo_simple_graph) 
op_repo = FakeOperatorRepository()
mapping_manual = [ParameterMapping(
    src_operator_graph_id="b342a8bd-21db-47be-89a1-0db7ab1798fb",
    source_operator_id="add",
    dst_operator_graph_id="9fbad8cb-110e-4cd3-97ee-cbe7919b8753",
    dst_operator_id="multiple",
    src_param_name="y",
    dst_param_name="e", 
)]
pms = ParameterMappingService(mapping_manual)
argv_for_execute = {"param": {"a": 1, "b": 4, "c":5, "d":6}}
a = Algorithm(algo_info, OperatorComponentToOperatorInfoAdapter(op_repo), pms, algo_graph, **argv_for_execute)
task_factory = AlgoTaskFactory(a)

tasks = build_algo_tasks(task_factory, a)
exec = DefaultExecutor()
res = exec.invoke_tasks(tasks)
logging.info("local calc results %s", res)

from calc_manager.ray.ray_manager import RayManager
from calc_manager.ray.remote_task import RayAlgoOperatorTaskFactory
import calc_manager
class RayFakeOperatorRepository(OperatorRepository):
    def __init__(self):
        sample_add = {
            "operator_name": "add",
            "operator_path": "sample_add/add.py",
            "operator_func_entrypoint": "add",
            "operator_desc": "This is a test operator",
        }
        sample_mul = {
            "operator_name": "multiple",
            "operator_path": "sample_multiple/mul.py",
            "operator_func_entrypoint": "mul",
            "operator_desc": "This is a test operator",
        }
        self.repo = {"add": sample_add, "multiple": sample_mul}
        

    def get_op_by_id(self, op_id: str) -> Optional[OperatorInfo]:
        info = self.repo[op_id]
        return OperatorInfo(**info)

manager = RayManager(
    "ray://10.130.9.3:10001",
    {
        "logging_level": logging.INFO,
        "runtime_env": {
            "py_modules": [calc_manager],
            "working_dir": "./operators_sample",
            # "pip": ["emoji", "numpy", "pynvim"],
        }
    },
)
op_repo = RayFakeOperatorRepository()
a = Algorithm(algo_info, OperatorComponentToOperatorInfoAdapter(op_repo), pms, algo_graph, **argv_for_execute)

ray_task_factory = RayAlgoOperatorTaskFactory(a)
ray_tasks = build_algo_tasks(ray_task_factory, a)
res_ray, ctx = manager.invoke_tasks(ray_tasks)
print(res_ray)